import pandas as pd
import numpy as np
import matplotlib.pyplot as plt
from scipy import stats


plt.rcParams["font.sans-serif"] = ["SimHei"]
plt.rcParams["axes.unicode_minus"] = False
# 加载数据
data = pd.read_csv('./boston.csv')

# 设置图片清晰度
plt.rcParams['figure.dpi'] = 300

# 将 AGE 分 5 个区间
data['AGE_bin'] = pd.qcut(data['AGE'], q=5)

# 计算各区间的 MEDV 均值、数量和 95% 置信区间
grouped = data.groupby('AGE_bin')['MEDV']
means = grouped.mean()
n = grouped.count()
std_err = grouped.sem()
ci = std_err * stats.t.ppf((1 + 0.95) / 2., n - 1)

# 创建图形和坐标轴
fig, ax1 = plt.subplots(figsize=(10, 6))

# 背景添加原始数据点的半透明显示
ax1.scatter(data['AGE'], data['MEDV'], alpha=0.2, color='gray', label='原始数据点')

# 点线图连接各区间均值，设置颜色为黑色
line, = ax1.plot(means.index.astype(str), means, marker='o', color='black', label='各区间均值')

# 设置误差棒颜色从冷到暖变化
cmap = plt.get_cmap('coolwarm')
colors = [cmap(i) for i in np.linspace(0, 1, len(means))]

# 绘制误差棒，添加水平标记线
for i, (mean, err, color) in enumerate(zip(means, ci, colors)):
    ax1.errorbar(i, mean, yerr=err, fmt='none', color=color, capsize=5, elinewidth=1, capthick=1)
    ax1.axhline(y=mean + err, xmin=i/len(means), xmax=(i + 1)/len(means), color=color, linestyle='--', linewidth=0.5)
    ax1.axhline(y=mean - err, xmin=i/len(means), xmax=(i + 1)/len(means), color=color, linestyle='--', linewidth=0.5)

# 设置主坐标轴标签和标题
ax1.set_xlabel('房龄区间')
ax1.set_ylabel('房价均值')
ax1.set_title('不同房龄区间的房价置信区间')

# 创建右侧副轴
ax2 = ax1.twinx()

# 在副轴上绘制各区间房屋数量
bars = ax2.bar(means.index.astype(str), n, width=0.3, alpha=0.3, color='orange', label='房屋数量')
ax2.set_ylabel('房屋数量')

# 合并图例
lines, labels = ax1.get_legend_handles_labels()
lines2, labels2 = ax2.get_legend_handles_labels()
ax2.legend(lines + lines2, labels + labels2, loc='upper left')

# 显示图形
plt.show()